package dyyx.echo;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufHolder;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelHandler.Sharable;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;

@Sharable
public class TraficCountHandler extends ChannelDuplexHandler {
	
	static final AtomicLong writeBytes = new AtomicLong(0);
	static final AtomicLong readBytes = new AtomicLong(0);


	@Override
	public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
		long bytes = calculateSize(msg);
		if(bytes>0){
			readBytes.addAndGet(bytes);
		}
		ctx.fireChannelRead(msg);
	}

	@Override
	public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {		
		long bytes = calculateSize(msg);
		if(bytes>0){
			writeBytes.addAndGet(bytes);
		}
		ctx.write(msg, promise);
	}
	
	private long calculateSize(Object msg) {
        if (msg instanceof ByteBuf) {
            return ((ByteBuf) msg).readableBytes();
        }
        if (msg instanceof ByteBufHolder) {
            return ((ByteBufHolder) msg).content().readableBytes();
        }
        return -1;
    }
	
	public static Map getInfo(){
		Map map = new HashMap();
		map.put("writeBytes", writeBytes);
		map.put("readBytes", readBytes);
		return map;
	}

}
